package com.gaogzhen.datastructure.tree;

import com.gaogzhen.datastructure.queue.LinkQueue;
import com.gaogzhen.datastructure.queue.Queue;
import com.gaogzhen.datastructure.stack.Stack;

import java.util.Iterator;
import java.util.NoSuchElementException;

/**
 * 非递归红黑二叉查找树
 * 2-3树的对应的一种实现
 * 定义：
 * 红链接均为左链接
 * 没有任何一个结点同时和两条红色链接相连
 * 该树是完美黑色平衡的，即任意空链接到根结点的路径上的黑链接数量相同
 * 与2-3树一一对应
 * 红链接将2两个2-结点连起来构成一个3-结点。确切的说，我们将3-结点表示为一条左斜的红链接链接的2个阶段
 * 两个2-结点其中一种是另外一个的左子结点
 * 黑链接则是2-3树中的普通链接即2-链接
 *
 * @author Administrator
 * @date 2022-12-15 20:43
 */
public class RedBlackTree<K extends Comparable<K>, V> {

    /**
     * 根结点
     */
    private Node root;

    /**
     * 链接（结点）颜色 RED = true表示红链接；BLACK = false表示黑链接
     */
    private static final boolean RED = true;
    private static final boolean BLACK = false;

    /**
     * 返回树中键值对数量
     *
     * @return 树中键值对数量
     */

    public int size() {
        return size(root);
    }

    /**
     * 红黑树结点key的迭代
     * 中序遍历红黑树
     *
     * @return 红黑树结点key的迭代器
     */

    public Iterable<K> keys() {
        if (isEmpty()) {
            return null;
        }
        // 栈存放遍历路径（结点）
        Stack<Node> s = new Stack<>();
        // 队列存放最终结果
        Queue<K> q = new LinkQueue<>();
        Node x = root;
        while (x != null || !s.isEmpty()) {
            while (x != null) {
                s.push(x);
                x = x.left;
            }
            if (!s.isEmpty()) {
                x = s.pop();
                q.offer(x.key);
                x = x.right;
            }

        }
        return q;
    }

    /**
     * 返回目标结点为根结点的子树键值对数量
     *
     * @param x 目标结点
     * @return 目标结点为根结点的子树键值对数量
     */
    private int size(Node x) {
        return x == null ? 0 : x.n;
    }

    /**
     * 判断红黑二叉树是否空
     *
     * @return 红黑二叉树为空返回 {@code true}；否则返回{@code false}
     */

    public boolean isEmpty() {
        return root == null;
    }

    /**
     * 插入键值对
     *
     * @param key   键
     * @param value 值
     */
    public void put(K key, V value) {
        // 树为空
        if (root == null) {
            root = new Node(key, value);
            return;
        }
        // 查找key是否在树中，且记录访问路径
        Stack<Node> stack = new Stack<>();
        Node cur = root;
        Node p = null;
        int type = 0;
        while (cur != null) {
            stack.push(cur);
            p = cur;
            int cmp = key.compareTo(cur.key);
            if (cmp < 0) {
                cur = cur.left;
                type = 1;
            } else if (cmp > 0) {
                cur = cur.right;
                type = 2;
            } else {
                cur.value = value;
                return;
            }
        }
        // key不在树中
        Node n = new Node(key, value);
        if (type == 1) {
            // 新结点为左子结点
            p.left = n;
        } else {
            // 新结点为右子结点
            p.right = n;
        }

        // 插入新结点后，为满足红黑树性质沿遍历路径自下向上做平衡调整
        while (!stack.isEmpty()) {
            cur = stack.pop();
            cur.n += 1;
            // 如果右链接为红色且左链接为黑色，则左旋
            if (isRed(cur.right) && !isRed(cur.left)) {
                if (cur == root) {
                    root = cur = rotateLeft(cur);
                } else {
                    p = stack.peek();
                    if (p.left == cur) {
                        p.left = cur = rotateLeft(cur);
                    } else {
                        p.right = cur = rotateLeft(cur);
                    }
                }
            }
            // 如果左链接和左链接的左链接都为红色，则右旋
            if (isRed(cur.left) && isRed(cur.left.left)) {
                if (cur == root) {
                    root = cur = rotateRight(cur);
                } else {
                    p = stack.peek();
                    if (p.left == cur) {
                        p.left = cur = rotateRight(cur);
                    } else {
                        p.right = cur = rotateRight(cur);
                    }
                }
            }
            // 如果左右链接都为红色，则把红色转移到父链接
            if (isRed(cur.left) && isRed(cur.right)) {
                flipColors(cur);
            }
        }
        // 保证根结点为黑色
        root.color = BLACK;
    }

    /**
     * 获取key对应的value值
     * 因红黑树为二叉查找树，效果等同于二分查找
     *
     * @param key 键
     * @return 返回key对应的value值
     */

    public V get(K key) {
        // 树为空
        if (root == null) {
            return null;
        }
        // 查找key是否在树中
        Node x = root;
        while (x != null) {
            int cmp = key.compareTo(x.key);
            if (cmp < 0) {
                // 目标key比当前key小，继续在左子树中查找
                x = x.left;
            } else if (cmp > 0) {
                // 目标key比当前key大，继续在右子树中查找
                x = x.right;
            } else {
                // 目标key与当前key相等，命中目标，返回对应的value值
                return x.value;
            }
        }
        // key不在树中，返回null
        return null;
    }

    /**
     * 删除指定的key
     *
     * @param key 指定key
     */
    public void delete(K key) {
        if (key == null) {
            throw new IllegalArgumentException("argument to delete() is null");
        }
        if (!contains(key)) {
            return;
        }

        // if both children of root are black, set root to red
        if (!isRed(root.left) && !isRed(root.right)) {
            root.color = RED;
        }

        root = delete(root, key);
        if (!isEmpty()) {
            root.color = BLACK;
        }
    }

    public boolean contains(K key) {
        return get(key) != null;
    }

    /**
     * 在以h为根结点的树中删除指定key
     *
     * @param h   树的根结点
     * @param key key
     * @return 删除后调整完的新树
     */
    private Node delete(Node h, K key) {
        // assert get(h, key) != null;
        // 存储遍历的结点
        Stack<Node> stack = new Stack<>();
        Node t;
        Node pp = null;
        while (true) {
            if (key.compareTo(h.key) < 0) {
                // 比当前结点key小，确保当前结点不为2-结点
                if (!isRed(h.left) && !isRed(h.left.left)) {
                    // 当前结点为2-结点，从兄弟结点借一个结点
                    t = moveRedLeft(h);
                    if (pp != null) {
                        adjustParentLink(h, pp, t);
                    }
                    h = t;
                }
                stack.push(h);
                // 继续遍历左子树
                pp = h;
                h = h.left;
            } else {
                if (isRed(h.left)) {
                    t = rotateRight(h);
                    if (pp != null) {
                        adjustParentLink(h, pp, t);
                    }
                    h = t;
                }
                if (key.compareTo(h.key) == 0 && (h.right == null)) {
                    // 命中叶子结点
                    if (stack.isEmpty()) {
                        // 要删除的为根结点
                        return null;
                    } else {
                        // 非根结点，父链接置空
                        Node p = stack.peek();
                        if (p.left == h) {
                            p.left = null;
                        } else {
                            p.right = null;
                        }
                    }
                    break;
                }
                if (!isRed(h.right) && !isRed(h.right.left)) {
                    // 右子树为2-结点，从左子树结点借
                    t = moveRedRight(h);
                    if (pp != null) {
                        adjustParentLink(h, pp, t);
                    }
                    h = t;
                }
                if (key.compareTo(h.key) == 0) {
                    // 命中非叶子结点，与后继结点交换
                    Node m = min(h.right);
                    swap(h, m);
                    // 删除替换后的后继结点
                    h.right = deleteMin(h.right);
                    break;
                } else {
                    // 继续遍历右子树
                    stack.push(h);
                    h = h.right;
                }
            }
        }
        // 调整删除结点后的二叉树，以满足红黑树的性质
        Node b;
        while (!stack.isEmpty()) {
            h = stack.pop();
            h.n--;
            b = balance(h);
            if (!stack.isEmpty()) {
                t = stack.peek();
                adjustParentLink(h, t, b);
            } else {
                return b;
            }
        }
        return null;
    }

    /**
     * 新结点替换旧结点根据旧结点是父结点的左子树，还是右子树的判断结果赋值
     *
     * @param old  旧结点
     * @param p    父结点
     * @param node 新结点
     */
    private void adjustParentLink(Node old, Node p, Node node) {
        if (p.left == old) {
            p.left = node;
        } else {
            p.right = node;
        }
    }

    /**
     * 交换两个结点的key和value
     *
     * @param h 结点
     * @param m 结点
     */
    private void swap(Node h, Node m) {
        K k = h.key;
        V v = h.value;
        h.key = m.key;
        h.value = m.value;
        m.key = k;
        m.value = v;
    }

    /**
     * 查找最小的key对应的值
     *
     * @return 最小key对应的值
     */
    public K min() {
        if (isEmpty()) {
            throw new NoSuchElementException("calls min() with empty symbol table");
        }
        return min(root).key;
    }

    /**
     * 查找x为根结点的树中最小key对应的结点
     *
     * @param x 根结点
     * @return 最小key对应的结点
     */
    private Node min(Node x) {
        while (x.left != null) {
            x = x.left;
        }
        return x;
    }

    /**
     * 查找最大key对应的值
     *
     * @return 最大key对应的value
     */
    public K max() {
        if (isEmpty()) {
            throw new NoSuchElementException("calls max() with empty symbol table");
        }
        return max(root).key;
    }

    /**
     * 查找以x为根结点的树中最大key对应的结点
     *
     * @param x 根结点
     * @return 最大key对应的结点
     */
    private Node max(Node x) {
        while (x.right != null) {
            x = x.right;
        }
        return x;
    }

    /**
     * 查找小于等于给定key的最大键
     *
     * @param key 指定key
     * @return 小于等于给定key的最大键
     */
    public K floor(K key) {
        if (key == null) {
            throw new IllegalArgumentException("argument to floor() is null");
        }
        if (isEmpty()) {
            throw new NoSuchElementException("calls floor() with empty symbol table");
        }
        Node x = floor(root, key);
        if (x == null) {
            throw new NoSuchElementException("argument to floor() is too small");
        } else {
            return x.key;
        }
    }

    /**
     * 查找以x为根结点树中小于等于给定key的最大键
     *
     * @param key 指定key
     * @return 小于等于给定key的最大键
     */
    private Node floor(Node x, K key) {
        Node t = null;
        while (x != null) {
            int cmp = key.compareTo(x.key);
            if (cmp == 0) {
                // 名字直接返回
                return x;
            } else if (cmp < 0) {
                // 给定键小于当前结点键，继续遍历左子结点
                x = x.left;
            } else {
                // 目前为止当前结点为小于等于key的最大结点
                t = x;
                x = x.right;
            }
        }
        // 返回最终小于等于给定key的最大键
        return t;
    }

    /**
     * 查找大于等于key的最小键
     *
     * @param key 指定key
     * @return 大于等于key的最小键
     */
    public K ceiling(K key) {
        if (key == null) {
            throw new IllegalArgumentException("argument to ceiling() is null");
        }
        if (isEmpty()) {
            throw new NoSuchElementException("calls ceiling() with empty symbol table");
        }
        Node x = ceiling(root, key);
        if (x == null) {
            throw new NoSuchElementException("argument to ceiling() is too small");
        } else {
            return x.key;
        }
    }

    /**
     * 查找x为根结点的树中大于等于key的最小键
     *
     * @param key 指定key
     * @return 大于等于key的最小键
     */
    private Node ceiling(Node x, K key) {
        Node t = null;
        while (x != null) {
            int cmp = key.compareTo(x.key);
            if (cmp == 0) {
                // 命中直接返回
                return x;
            } else if (cmp > 0) {
                // 给定键大于key，继续遍历右子结点
                x = x.right;
            } else {
                // t记录当前大于等于key的最小键x，继续遍历左子结点
                t = x;
                x = x.left;
            }
        }
        // 返回最终大于等于key的最小键
        return t;
    }


    /**
     * 查找小于key的键的数量
     *
     * @param key 给定的key
     * @return 小于key的键的数量
     */
    public int rank(K key) {
        if (key == null) {
            throw new IllegalArgumentException("argument to rank() is null");
        }
        return rank(key, root);
    }

    /**
     * 查找以x为根结点小于key的键的数量
     * 非递归算法：
     * 初始计数t=0
     * 循环开始，当前结点不为空
     * 比较key与当前结点key大小
     * 如果key小于当前结点的key，继续查找左子树中的小于key的键的数量
     * 如果key大于当前结点的key，计数+1+当前结点左子树键数量;继续查找右子树中小于key的数量
     * 1 为当前结点
     * 如果key等于当前节点中的key，计数+当前结点左子树键的数量;循环结束
     * 返回计数
     *
     * @param key 给定的key
     * @return 小于key的键的数量
     */
    private int rank(K key, Node x) {
        int t = 0;
        while (x != null) {
            int cmp = key.compareTo(x.key);
            if (cmp < 0) {
                x = x.left;
            } else if (cmp > 0) {
                t += 1 + size(x.left);
                x = x.right;
            } else {
                t += size(x.left);
                break;
            }
        }
        return t;
    }


    /**
     * 排名(索引)为rank的键
     *
     * @param rank 指定排名
     * @return 排名(索引)为rank的键
     */
    public K select(int rank) {
        if (rank < 0 || rank >= size()) {
            throw new IllegalArgumentException("argument to select() is invalid: " + rank);
        }
        return select(root, rank);
    }

    /**
     * 以x为根结点树排名(索引)为rank的键
     *
     * @param rank 指定排名
     * @return 排名(索引)为rank的键
     */
    private K select(Node x, int rank) {
        int c = rank;
        while (x != null) {
            int leftSize = size(x.left);
            if (leftSize > c) {
                //
                x = x.left;
            } else if (leftSize < c) {
                c -= leftSize + 1;
                x = x.right;
            } else {
                return x.key;
            }
        }
        return null;
    }

    public void delete(int index) {

    }

    /**
     * 删除最小结点
     *
     * @return 最小结点对应的value
     */

    public void deleteMin() {
        if (isEmpty()) {
            throw new NoSuchElementException("BST underflow");
        }

        // 如果根结点左右链接都为黑色，变为4-结点
        if (!isRed(root.left) && !isRed(root.right)) {
            root.color = RED;
        }

        root = deleteMin(root);
        if (!isEmpty()) {
            root.color = BLACK;
        }
    }

    /**
     * 删除以p为根结点的树中最小结点
     *
     * @param p 根结点
     * @return 删除最小结点后的新树
     */
    private Node deleteMin(Node p) {
        // 目标结点的左子树
        Node pl;
        // 目标结点的父结点
        Node pp = null;
        // 栈记录遍历结点（路径），用于删除后自下向上调整红黑树，以满足红黑树的性质
        Stack<Node> stack = new Stack<>();
        while ((pl = p.left) != null) {
            // 如果左链接为2-结点，从兄弟结点或者父结点借
            if (!isRed(pl) && !isRed(pl.left)) {
                p = moveRedLeft(p);
                if (pp != null) {
                    pp.left = p;
                }
            }
            // 记录遍历结点
            stack.push(p);
            // 继续沿左子树遍历
            pp = p;
            p = pp.left;
        }

        // pp为空说明要删除的为根结点
        if (pp == null) {
            return null;
        } else {
            // 要删除的结点为非根结点，父结点的左链接置空
            pp.left = null;
        }
        // 调整删除结点后的二叉树，以满足红黑树的性质
        while (!stack.isEmpty()) {
            p = stack.pop();
            p.n--;
            if (!stack.isEmpty()) {
                stack.peek().left = balance(p);
            } else {
                return balance(p);
            }
        }
        return null;
    }

    /**
     * 平衡调整，使修改后的子树保持红黑树的特征
     *
     * @param h 目标子树根结点
     * @return 调整后的子树根结点
     */
    private Node balance(Node h) {
        // assert (h != null);

        // 结点右链接为红色且左链接为黑色，左旋
        if (isRed(h.right) && !isRed(h.left)) {
            h = rotateLeft(h);
        }

        // 结点左链接为红色且左链接的左链接也为红色，右旋
        if (isRed(h.left) && isRed(h.left.left)) {
            h = rotateRight(h);
        }

        // 左链接和右链接都为红色，左右链接置为黑色，红色转义到父链接
        if (isRed(h.left) && isRed(h.right)) {
            flipColors(h);
        }
        return h;
    }


    /**
     * 如果h是红结点，h.left和h.left.left是黑结点，确保h.left或者h.left的子结点其中之一为红结点
     *
     * @param h 根结点
     * @return 调整后的新树
     */
    private Node moveRedLeft(Node h) {
        // assert (h != null);
        // assert isRed(h) && !isRed(h.left) && !isRed(h.left.left);

        flipColors(h);
        if (isRed(h.right.left)) {
            h.right = rotateRight(h.right);
            h = rotateLeft(h);
            flipColors(h);
        }
        return h;
    }

    /**
     * 如果h是红结点，h.right和h.right.left是黑结点，确保h.right或者h.right的子结点其中之一为红结点
     *
     * @param h 根结点
     * @return 调整后的新树
     */
    private Node moveRedRight(Node h) {
        // assert (h != null);
        // assert isRed(h) && !isRed(h.right) && !isRed(h.right.left);
        flipColors(h);
        if (isRed(h.left.left)) {
            h = rotateRight(h);
            flipColors(h);
        }
        return h;
    }

    /**
     * 删除最大结点
     */
    public void deleteMax() {
        if (isEmpty()) {
            throw new NoSuchElementException("BST underflow");
        }

        // 如果根结点左右链接都为黑色，变为4-结点
        if (!isRed(root.left) && !isRed(root.right)) {
            root.color = RED;
        }
        root = deleteMax(root);
        if (!isEmpty()) {
            root.color = BLACK;
        }
    }

    /**
     * 删除p为根结点的最大结点
     *
     * @param p 根结点
     * @return 删除最大结点后的新树
     */
    private Node deleteMax(Node p) {
        // 目标结点的右子树
        Node pr;
        // 目标结点的父结点
        Node pp = null;
        // 栈记录遍历结点（路径），用于删除后自下向上调整红黑树，以满足红黑树的性质
        Stack<Node> stack = new Stack<>();
        while (true) {
            // 若左子树为红链接，右旋，保证右子树为3-结点
            if (isRed(p.left)) {
                p = rotateRight(p);
                if (pp != null) {
                    pp.right = p;
                }
            }
            if ((pr = p.right) == null) {
                break;
            }
            // 如果右链接为2-结点，从兄弟结点或者父结点借
            if (!isRed(pr) && !isRed(pr.left)) {
                p = moveRedRight(p);
                if (pp != null) {
                    pp.right = p;
                }
            }
            // 记录遍历结点
            stack.push(p);
            // 继续沿右子树遍历
            pp = p;
            p = pp.right;
        }
        // pp为空说明要删除的为根结点
        if (pp == null) {
            return null;
        } else {
            // 要删除的结点为非根结点，父结点的右链接置空
            pp.right = null;
        }
        // 调整删除结点后的二叉树，以满足红黑树的性质
        while (!stack.isEmpty()) {
            p = stack.pop();
            p.n--;
            if (!stack.isEmpty()) {
                stack.peek().right = balance(p);
            } else {
                return balance(p);
            }
        }
        return null;
    }

    /**
     * [lo,hi]之间键的数量
     *
     * @param lo 低位key
     * @param hi 高位key
     * @return [lo, hi]之间键的数量
     */
    public int size(K lo, K hi) {
        if (lo == null) {
            throw new IllegalArgumentException("first argument to size() is null");
        }
        if (hi == null) {
            throw new IllegalArgumentException("second argument to size() is null");
        }

        if (lo.compareTo(hi) > 0) {
            return 0;
        }
        if (contains(hi)) {
            return rank(hi) - rank(lo) + 1;
        } else {
            return rank(hi) - rank(lo);
        }
    }

    /**
     * 位于[lo,hi]之间键的集合
     *
     * @param lo 低位key
     * @param hi 高位key
     * @return [lo, hi]之间键的集合
     */
    public Iterable<K> keys(K lo, K hi) {
        if (lo == null) {
            throw new IllegalArgumentException("first argument to keys() is null");
        }
        if (hi == null) {
            throw new IllegalArgumentException("second argument to keys() is null");
        }

        Queue<K> queue = new LinkQueue<>();
        // if (isEmpty() || lo.compareTo(hi) > 0) return queue;
        keys(root, queue, lo, hi);
        return queue;
    }

    /**
     * 以x为根结点树中位于[lo,hi]之间键的集合
     *
     * @param lo 低位key
     * @param hi 高位key
     * @return [lo, hi]之间键的集合
     */
    private void keys(Node x, Queue<K> queue, K lo, K hi) {

        Stack<Node> s = new Stack<>();
        while (x != null || !s.isEmpty()) {
            while (x != null && lo.compareTo(x.key) < 0) {
                // 当前结点键大于低位键lo，继续遍历左子树
                s.push(x);
                x = x.left;
            }

            if (!s.isEmpty()) {
                if (x == null) {
                    x = s.pop();
                }
                // lo <= x.key <= hi ,加入队列
                if (lo.compareTo(x.key) <= 0 && hi.compareTo(x.key) >= 0) {
                    queue.offer(x.key);
                }
            }
            if (x != null && hi.compareTo(x.key) > 0) {
                // 当前结点key小于最大值hi,继续遍历右子树
                x = x.right;
            } else {
                // 当前结点key大于等于最大值hi，结束
                break;
            }

        }
    }

    /**
     * 结点类
     */
    final class Node {
        /**
         * 键
         */
        K key;

        /**
         * 值
         */
        V value;

        /**
         * 左右链接（子结点）
         */
        Node left, right;

        /**
         * 子树结点个数
         */
        int n;

        /**
         * 链接颜色
         */
        boolean color;

        public Node(K key, V value, Node left, Node right) {
            this.key = key;
            this.value = value;
            this.left = left;
            this.right = right;
            // 默认新建结点为根结点的子树结点个数为1
            this.n = 1;
            // 默认新建结点颜色为红色
            this.color = RED;
        }

        public Node(K key, V value) {
            this(key, value, null, null);
        }
    }

    /**
     * 判断当前结点和父链接是否为红色
     *
     * @param x 当前结点
     * @return 当前结点不为空且当前结点color属性为RED，返回{@code true}；否则返回{@code false}
     */
    private boolean isRed(Node x) {
        return x != null && x.color == RED;
    }

    /**
     * 以子树根结点为轴左旋
     *
     * @param p 子树根结点
     * @return 左旋完成后的子树根结点
     */
    private Node rotateLeft(Node p) {
        // 调整链接
        Node x = p.right;
        int n = 1 + (x.right == null ? 0 : x.right.n);
        p.right = x.left;
        x.left = p;
        // 调整颜色
        x.color = p.color;
        p.color = RED;
        // 计算结点数
        x.n = p.n;
        p.n = x.n - n;
        return x;
    }

    /**
     * 以子树根结点为轴右旋
     *
     * @param p 子树根结点
     * @return 右旋完成后的子树根结点
     */
    private Node rotateRight(Node p) {
        // 调整链接
        Node x = p.left;
        int n = 1 + (x.left == null ? 0 : x.left.n);
        p.left = x.right;
        x.right = p;
        // 调整颜色
        x.color = p.color;
        p.color = RED;
        // 计算结点数
        x.n = p.n;
        p.n = x.n - n;
        return x;
    }

    /**
     * 变化双子链接为红色为黑色，红色转义到父结点
     *
     * @param p 当前结点
     */
    private void flipColors(Node p) {
        p.color = !p.color;
        p.left.color = !p.left.color;
        p.right.color = !p.right.color;
    }

    @Override
    public String toString() {
        Iterable<K> iterable = keys();
        Iterator<K> iterator;
        if (iterable == null || !((iterator = iterable.iterator()).hasNext())) {
            return "{}";
        }
        StringBuilder sb = new StringBuilder();
        sb.append("{");
        while (true) {
            K k = iterator.next();
            V v = get(k);
            sb.append(k).append("=").append(v);
            if (!iterator.hasNext()) {
                sb.append("}");
                return sb.toString();
            }
            sb.append(", ");
        }
    }
}
